"""数据库连接模块"""
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError
import sys


class DatabaseConnector:
    """数据库连接器"""
    
    def __init__(self, host, port, database, db_type, user, password):
        """
        初始化数据库连接器
        
        Args:
            host: 数据库主机地址
            port: 数据库端口
            database: 数据库名称
            db_type: 数据库类型 ('mysql' 或 'postgresql')
            user: 用户名
            password: 密码
        """
        self.host = host
        self.port = port
        self.database = database
        self.db_type = db_type
        self.user = user
        self.password = password
        self.engine = None
        
    def connect(self):
        """
        建立数据库连接
        
        Returns:
            engine: SQLAlchemy engine 对象
            
        Raises:
            Exception: 连接失败时抛出异常
        """
        try:
            # 构建连接字符串
            if self.db_type == 'mysql':
                connection_string = (
                    f"mysql+pymysql://{self.user}:{self.password}"
                    f"@{self.host}:{self.port}/{self.database}"
                    f"?charset=utf8mb4"
                )
            elif self.db_type == 'postgresql':
                connection_string = (
                    f"postgresql+psycopg2://{self.user}:{self.password}"
                    f"@{self.host}:{self.port}/{self.database}"
                )
            else:
                raise ValueError(f"不支持的数据库类型: {self.db_type}")
            
            # 创建引擎
            engine_kwargs = {
                'pool_pre_ping': True,  # 连接前测试连接
                'echo': False
            }
            
            # 尝试创建引擎并测试连接
            try:
                self.engine = create_engine(connection_string, **engine_kwargs)
                # 测试连接
                with self.engine.connect() as conn:
                    conn.execute(text("SELECT 1"))
            except (SQLAlchemyError, Exception) as version_error:
                # 如果出现版本解析错误（通常是 KingbaseES 等兼容 PostgreSQL 的数据库）
                error_msg = str(version_error.orig) if hasattr(version_error, 'orig') else str(version_error)
                if "Could not determine version" in error_msg or "KingbaseES" in error_msg:
                    # 对于 KingbaseES 等数据库，手动设置一个兼容的 PostgreSQL 版本号
                    # 使用 PostgreSQL 14.0 版本号来避免版本解析错误
                    engine_kwargs['server_version_info'] = (14, 0)
                    try:
                        self.engine = create_engine(connection_string, **engine_kwargs)
                        # 再次测试连接
                        with self.engine.connect() as conn:
                            conn.execute(text("SELECT 1"))
                    except Exception as retry_error:
                        # 如果重试仍然失败，抛出原始错误或重试错误
                        raise version_error
                else:
                    # 其他错误直接抛出，让外层异常处理处理
                    raise
            
            return self.engine
            
        except SQLAlchemyError as e:
            error_msg = str(e.orig) if hasattr(e, 'orig') else str(e)
            if "Access denied" in error_msg or "authentication failed" in error_msg.lower():
                raise Exception(f"数据库认证失败，请检查用户名和密码: {error_msg}")
            elif "Unknown database" in error_msg or "database" in error_msg.lower() and "does not exist" in error_msg.lower():
                raise Exception(f"数据库 '{self.database}' 不存在: {error_msg}")
            elif "Can't connect" in error_msg or "could not connect" in error_msg.lower():
                raise Exception(f"无法连接到数据库服务器 ({self.host}:{self.port})，请检查 host、port 和网络连接: {error_msg}")
            else:
                raise Exception(f"数据库连接失败: {error_msg}")
        except Exception as e:
            raise Exception(f"数据库连接失败: {str(e)}")
    
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()

